#!/usr/bin/python

class ChecksummingFilehandle:
    # fh can be something like open('/tmp/blah', 'rb')
    # hasher can be something like hashlib.md5()
    def __init__(self, fh, hasher):
        self.fh = fh
        self.hasher = hasher
        self.bytes_read = 0
    def read(self, n):
        ret = self.fh.read(n)
        self.bytes_read += len(ret)
        self.hasher.update(ret)
        return ret
    def close(self):
        return self.fh.close()
    def __enter__(self):
        return self
    def __exit__(self, type, value, traceback):
        self.close()

# Only allow reads of the next 'limit' bytes of 'fh'.
class LimitedFilehandle:
    def __init__(self, fh, limit):
        self.fh = fh
        self.remaining_bytes = limit
    def read(self, n):
        ret = self.fh.read(min(n, self.remaining_bytes))
        self.remaining_bytes -= len(ret)
        return ret
    def close(self):
        return self.fh.close()
    def __enter__(self):
        return self
    def __exit__(self, type, value, traceback):
        self.close()
